from flask import request
from flask_socketio import send, Namespace
from datetime import datetime
from .share import shared_vars

# 定义 WebSocket 命名空间
class WorkerWebSocketNamespace(Namespace):
    def on_connect(self):
        print('Client connected')
        # 获取当前客户端的唯一连接标识符
        client_sid = request.sid
        print(f"客户端已连接，SID: {client_sid}")
        # 可以在这里发送初始数据或任务，room指定只发给自己
        self.emit('message', 'Welcome to the ACE mvp prototype server!', room=client_sid)

    def on_disconnect(self):
        print('client disconnected')
        client_sid = request.sid
        print(f"客户端已连接，SID: {client_sid}")
        shared_vars['workers'] = [node for node in shared_vars['workers'] if not node['sid']==client_sid ]
        pass

    def on_register(self, capabilities):
        client_sid = request.sid
        print(f"on_register 客户端已连接，SID: {client_sid}")
        node = {
                'sid'              : client_sid,
                'worker_id'        : capabilities['worker_id'],
                'os_version'       : capabilities['os'],
                'boot_time'        : datetime.now().strftime('%Y-%m-%d %H : %M : %S'),
                'available_cpu'    : capabilities['cpus'],
                'available_memory' : capabilities['memory'],
                'running_tasks'    : -1,
                'version'          : capabilities['version'],
                'is_enabled'       : 1,
        }
        shared_vars['workers'].append(node)

    def on_task_finish(self, task):
        print('Task received:',task)
        # 处理任务并返回结果
        self.emit('task_result', {'task': task,'result':'任务已完成'})
